package com.itheima.algorithm.btree;

import java.util.Arrays;

/**
 * @author: TylerZhong
 * @description:
 */
public class BTree {
    public static class Node {
        public int[] keys; // 关键字
        public Node[] children; // 孩子
        public int keyNumber; // 有效关键字数目
        public boolean leaf = true; // 是否是叶子节点
        public int t; // 最小度数（最小孩子数）

        public Node(int t) { // t>=2
            this.t = t;
            this.children = new Node[2 * t];
            this.keys = new int[2 * t - 1];
        }

        public Node(int[] keys) {
            this.keys = keys;
        }

        @Override
        public String toString() {
            return Arrays.toString(Arrays.copyOfRange(keys, 0, keyNumber));
        }

        // 多路查找
        public Node get(int key) {
            int i = 0;
            while (i < keyNumber) {
                if (key == keys[i]) {
                    return this;
                }
                if (keys[i] > key) {
                    break;
                }
                i++;
            }
            // 执行到此时 keys[i] > key 或 i==keyNumber
            if (leaf) {
                return null;
            }
            // 非叶子情况
            return children[i].get(key);
        }

        // 向指定索引处插入 key
        public void insertKey(int key, int index) {
            System.arraycopy(keys, index, keys, index+1, keyNumber - index);
            keys[index] = key;
            keyNumber++;
        }

        // 向指定索引处插入 child
        public void insertChild(Node child, int index) {
            // 从后往前移动元素
            System.arraycopy(children, index, children, index + 1, keyNumber - index);
            children[index] = child;
        }

        // 移除指定 index 处的 key
        int removeKey(int index){
            int t = keys[index];
            System.arraycopy(keys, index+1, keys, index, --keyNumber - index);
            return t;
        }

        // 移除最左边的 key
        int removeLeftmostKey(){
            return removeKey(0);
        }

        // 移除最右边的 key
        int removeRightmostKey() {
            return removeKey(keyNumber - 1);
        }

        // 移除指定 index 处的 child
        Node removeChild(int index) {
            Node t = children[index];
            System.arraycopy(children, index + 1, children, index, keyNumber - index);
            children[keyNumber] = null; // help GC
            return t;
        }

        // 移除最左边的 child
        Node removeLeftmostChild() {
            return removeChild(0);
        }

        // 移除最右边的 child
        Node removeRightmostChild() {
            return removeChild(keyNumber);
        }

        // index 孩子处左边的兄弟
        Node childLeftSibling(int index) {
            return index > 0 ? children[index - 1] : null;
        }

        // index 孩子处右边的兄弟
        Node childRightSibling(int index) {
            return index == keyNumber ? null : children[index + 1];
        }

        // 复制当前节点的所有 key 和 child 到 target
        void moveToTarget(Node target) {
            int start = target.keyNumber;
            if (!leaf) {
                for (int i = 0; i <= keyNumber; i++) {
                    target.children[start + i] = children[i];
                }
            }

            for (int i = 0; i < keyNumber; i++) {
                target.keys[target.keyNumber++] = keys[i];
            }
        }

    }

    public Node root;
    int t; // 树中节点最小度数
    final int MIN_KEY_NUMBER; // 最小key 数目
    final int MAX_KEY_NUMBER; // 最大 key 数目

    public BTree() {
        this.t = 2;
        root = new Node(t);
        this.MAX_KEY_NUMBER = 2 * this.t - 1;
        this.MIN_KEY_NUMBER = this.t - 1;
    }

    public BTree(int t) {
        this.t = t;
        root = new Node(t);
        this.MAX_KEY_NUMBER = 2 * t - 1;
        this.MIN_KEY_NUMBER = t - 1;
    }

    // 1. 是否存在
    public boolean contains(int key) {
        return root.get(key) != null;
    }

    // 2. 新增
    public void put(int key) {
        doPut(root, key, null, 0);
    }

    private void doPut(Node node, int key, Node parent, int index) {
        int i = 0;
        while (i < node.keyNumber) {
            if (key == node.keys[i]) {
                return; // 更新
            }
            if (node.keys[i] > key) {
                break; // 找到了插入位置，即为此时的 i
            }
            i++;
        }

        if (node.leaf) {
            node.insertKey(key, i);
        } else {
            doPut(node.children[i], key, node, i);
        }
        if (node.keyNumber == MAX_KEY_NUMBER) {
            split(node, parent, index);
        }
    }

    /**
     * 分裂方法
     * @param left - 要分裂的节点
     * @param parent - 分裂节点的父节点
     * @param index - 分裂节点是第几个孩子
     */
    public void split(Node left, Node parent, int index) {
        // 分裂的是根节点
        if (parent == null) {
            Node newRoot = new Node(t);
            newRoot.leaf = false;
            newRoot.insertChild(left, 0);
            this.root = newRoot;
            parent = newRoot;
        }
        // 1. 创建 right 节点，把 left 中 t 之后的key和child移动过去
        Node right = new Node(t);
        right.leaf = left.leaf;
        System.arraycopy(left.keys, t, right.keys, 0, t - 1);
        if (!left.leaf) {
            System.arraycopy(left.children, t, right.children, 0, t);
        }
        right.keyNumber = t-1;
        left.keyNumber = t-1;
        // 2. 中间的 key (t-1处）插入到父节点
        int mid = left.keys[t - 1];
        parent.insertKey(mid, index);
        // 3. right 节点作为父节点的孩子
        parent.insertChild(right, index+1);
    }

    // 3. 删除
    public void remove(int key) {
        doRemove(null, root, 0, key);
    }

    private void doRemove(Node parent, Node node, int index, int key){
        int i = 0;
        while (i < node.keyNumber) {
            if (node.keys[i] >= key) {
                break;
            }
            i++;
        }
        // i 找到： 代表待删除 key 的索引
        // i 没找到：代表到第i个孩子继续查找
        if (node.leaf) {
            if (!found(node, key, i)) {
                // case 1
                return;
            } else {
                // case 2
                node.removeKey(i);
            }
        } else {
            if (!found(node, key, i)) {
                // case 3
                doRemove(node, node.children[i], i , key);
            } else {
                // case 4
                // 1. 找到后继 key
                Node s = node.children[i + 1];
                while (!s.leaf) {
                    s = s.children[0];
                }
                node.keys[i] = s.keys[0];
                // 3. 删除后继 Key
                doRemove(node, node.children[i+1], i+1, s.keys[0]);
            }
        }
        if (node.keyNumber < MIN_KEY_NUMBER) {
            // 调整平衡 case 5 6
            balance(parent, node, index);
        }
    }

    private boolean found(Node node, int key, int i) {
        return i < node.keyNumber && node.keys[i] == key;
    }

    private void balance(Node parent, Node x, int i) {
        // case 6 根节点
        if (x == root) {
            if (root.keyNumber == 0 && root.children[0] != null) {
                root = root.children[0];
            }
            return;
        }
        Node left = parent.childLeftSibling(i);
        Node right = parent.childRightSibling(i);
        if (left != null && left.keyNumber > MIN_KEY_NUMBER) {
            // case 5-1 左边富裕，右旋
            // a) 父亲节点中前驱 key 旋转下来
            x.insertKey(parent.keys[i - 1], 0);
            if (!left.leaf) {
                // b) left 中最大的孩子换爹
                x.insertChild(left.removeRightmostChild(), 0);
            }
            // c) left 中最大的key旋转上去
            parent.keys[i-1]= left.removeRightmostKey();
            return;
        }
        // case 5-2 右边富裕，左旋
        if (right != null && right.keyNumber > MIN_KEY_NUMBER) {
            // a) 父亲节点中后继 key 旋转下来
            x.insertKey(parent.keys[i], x.keyNumber);
            // b) right 中最小的孩子换爹
            if (!right.leaf) {
                x.insertChild(right.removeLeftmostChild(), x.keyNumber);
            }
            // c) right 中最小的 key 旋转上去
            parent.keys[i] = right.removeLeftmostKey();
            return;
        }
        if (left != null) {
            parent.removeChild(i);
            left.insertKey(parent.removeKey(i - 1), left.keyNumber);
            x.moveToTarget(left);
        } else {
            // 向自己合并
            parent.removeChild(i+1);
            x.insertKey(parent.removeKey(i), x.keyNumber);
            right.moveToTarget(x);
        }
    }

}
